#pragma once
#ifndef PWCONV_H_
#define PWCONV_H_

#include "stream_tools.h"

void MAC_s(ap_int<8> mul_a1, ap_int<8> mul_a2, ap_int<8> mul_b, ap_int<16>& rel1, ap_int<16>& rel2){
#pragma HLS INLINE
	ap_int<25> mul_a = ((ap_int<25>)mul_a1 << 16) + mul_a2;
	ap_int<32> result = mul_a * mul_b;
	rel1 = result(31,16) + result(15,15);
	rel2 = result(15,0);
}

void MAC_u(ap_int<8> mul_a1, ap_int<8> mul_a2, ap_uint<8> mul_b, ap_int<16>& rel1, ap_int<16>& rel2){
#pragma HLS INLINE
	ap_int<25> mul_a = ((ap_int<25>)mul_a1 << 16) + mul_a2;
	ap_int<32> result = mul_a * mul_b;
	rel1 = result(31,16) + result(15,15);
	rel2 = result(15,0);
}

template<
	unsigned N_IN,	  // # of INput
	unsigned N_OUT,	 // # of OUTput
	unsigned N_ICH,	 // # of Input CHannel
	unsigned N_OCH,	 // # of Output CHannel
	unsigned VEC_LEN,
	unsigned RAM = 1
>
void pwconv_u(
	data_stream<N_IN * BIT_ACTV>& in, 
	data_stream<N_OUT * BIT_CONV>& out, 
	const ap_uint<N_OUT * N_IN * BIT_WGHT> weight[N_OCH / N_OUT][N_ICH / N_IN]
)
{
	static_assert(N_ICH >= N_IN,  "pwconv_u");
	static_assert(N_OCH >= N_OUT, "pwconv_u");
	static_assert(N_ICH % N_IN  == 0, "pwconv_u");
	static_assert(N_OCH % N_OUT == 0, "pwconv_u");

	constexpr unsigned FOLD_I = N_ICH / N_IN;
	constexpr unsigned FOLD_O = N_OCH / N_OUT;
	constexpr unsigned ITERS = VEC_LEN;

	assert(in.size() == VEC_LEN * FOLD_I);
	assert(out.empty());

	if(RAM==0)
	{
		#pragma HLS bind_storage variable=weight type=rom_1p impl=lutram
	}
	else if(RAM==1)
	{
		#pragma HLS bind_storage variable=weight type=rom_1p impl=bram
	}
	else
	{
		#pragma HLS bind_storage variable=weight type=ram_1p impl=uram
	}

	ap_uint<N_IN * BIT_ACTV> line[FOLD_I];
	ap_int<BIT_CONV> acc[N_OUT];
	#pragma HLS array_partition variable=acc complete dim=0

	for (unsigned o = 0; o < N_OUT; ++o)
	{
		#pragma HLS UNROLL
		acc[o] = 0;
	}

	iters: for (unsigned it = 0; it < ITERS; ++it)
	{
		fold_o: for (unsigned fo = 0; fo < FOLD_O; ++fo)
		{
			fold_i: for (unsigned fi = 0; fi < FOLD_I; ++fi)
			{
				#pragma HLS PIPELINE II=1
				// load
				ap_uint<N_IN * BIT_ACTV> in_buf;
				if (fo == 0)
				{
					in_buf = in.read();
					line[fi] = in_buf;
				}
				else
				{
					in_buf = line[fi];
				}
				ap_uint<N_OUT * N_IN * BIT_WGHT> wt_buf = weight[fo][fi];

				for (unsigned i = 0; i < N_IN; ++i)
				{
				#pragma HLS UNROLL
					ap_uint<BIT_ACTV> x = in_buf(SLICE(BIT_ACTV, i));
					for (unsigned o = 0; o < N_OUT / 2; ++o)
					{
						ap_int<BIT_WGHT> w1 = wt_buf(SLICE(BIT_WGHT, N_IN * o * 2 + i));
						ap_int<BIT_WGHT> w2 = wt_buf(SLICE(BIT_WGHT, N_IN * o * 2 + N_IN + i));
						ap_int<16> tem1,tem2;
						MAC_u(w1, w2, x, tem1, tem2);
						acc[2 * o] += tem1;
						acc[2 * o + 1] += tem2;
					}
				}

				if(fi == FOLD_I-1)
				{
					ap_uint<N_OUT * BIT_CONV> out_buf;
					for (unsigned o = 0; o < N_OUT; ++o)
					{
						#pragma HLS UNROLL
						out_buf(SLICE(BIT_CONV, o)) = acc[o];
						acc[o] = 0;
					}
					out.write(out_buf);
				}
			}
		}
	}

	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD_O);
	return;
};

template<
	unsigned N_IN,	// # of INput
	unsigned N_OUT,	// # of OUTput
	unsigned N_ICH,	// # of Input CHannel
	unsigned N_OCH,	// # of Output CHannel
	unsigned VEC_LEN,
	unsigned RAM = 1
>
void pwconv_s(
	data_stream<N_IN * BIT_ACTV>& in,
	data_stream<N_OUT * BIT_CONV>& out,
	const ap_uint<N_OUT * N_IN * BIT_WGHT> weight[N_OCH / N_OUT][N_ICH / N_IN]
)
{
	static_assert(N_ICH >= N_IN,  "pwconv_s");
	static_assert(N_OCH >= N_OUT, "pwconv_s");
	static_assert(N_ICH % N_IN  == 0, "pwconv_s");
	static_assert(N_OCH % N_OUT == 0, "pwconv_s");

	constexpr unsigned FOLD_I = N_ICH / N_IN;
	constexpr unsigned FOLD_O = N_OCH / N_OUT;
	constexpr unsigned ITERS = VEC_LEN;

	assert(in.size() == VEC_LEN * FOLD_I);
	assert(out.empty());

	if(RAM==0)
	{
		#pragma HLS bind_storage variable=weight type=rom_1p impl=lutram
	}
	else if(RAM==1)
	{
		#pragma HLS bind_storage variable=weight type=rom_1p impl=bram
	}
	else
	{
		#pragma HLS bind_storage variable=weight type=ram_1p impl=uram
	}

	ap_uint<N_IN * BIT_ACTV> line[FOLD_I];
	ap_int<BIT_CONV> acc[N_OUT];
	#pragma HLS array_partition variable=acc complete dim=0

	for (unsigned o = 0; o < N_OUT; ++o)
	{
		#pragma HLS UNROLL
		acc[o] = 0;
	}

	iters: for (unsigned it = 0; it < ITERS; ++it)
	{
		fold_o: for (unsigned fo = 0; fo < FOLD_O; ++fo)
		{
			fold_i: for (unsigned fi = 0; fi < FOLD_I; ++fi)
			{
				#pragma HLS PIPELINE II=1
				// load
				ap_uint<N_IN * BIT_ACTV> in_buf;
				if (fo == 0)
				{
					in_buf = in.read();
					line[fi] = in_buf;
				}
				else
				{
					in_buf = line[fi];
				}
				ap_uint<N_OUT * N_IN * BIT_WGHT> wt_buf = weight[fo][fi];

				for (unsigned i = 0; i < N_IN; ++i)
				{
				#pragma HLS UNROLL
					ap_int<BIT_ACTV> x = in_buf(SLICE(BIT_ACTV, i));
					for (unsigned o = 0; o < N_OUT / 2; ++o)
					{
						ap_int<BIT_WGHT> w1 = wt_buf(SLICE(BIT_WGHT, N_IN * o * 2 + i));
						ap_int<BIT_WGHT> w2 = wt_buf(SLICE(BIT_WGHT, N_IN * o * 2 + N_IN + i));
						ap_int<16> tem1,tem2;
						MAC_s(w1, w2, x, tem1, tem2);
						acc[2 * o] += tem1;
						acc[2 * o + 1] += tem2;
					}
				}

				if(fi == FOLD_I-1)
				{
					ap_uint<N_OUT * BIT_CONV> out_buf;
					for (unsigned o = 0; o < N_OUT; ++o)
					{
						#pragma HLS UNROLL
						out_buf(SLICE(BIT_CONV, o)) = acc[o];
						acc[o] = 0;
					}
					out.write(out_buf);
				}
			}
		}
	}

	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD_O);
	return;
};

#endif
